# -*- coding: utf-8 -*-
"""
功能:
    对 JustRAIGS 数据集进行预处理。
    数据结构如下:
        JustRAIGS/
        ├── 0/
        ├── 1/
        ├── 2/
        ├── 3/
        ├── 4/
        ├── 5/
        └── JustRAIGS_Train_labels.csv

    说明:
        - 每个子目录 (0~5) 中存放 JPG 图像。
        - CSV 文件包含列:
          "Eye ID;Final Label;Fellow Eye ID;Age;Label G1;Label G2;Label G3;..."
        - 其中:
            * Eye ID → 图片文件名 (如 TRAIN000000 → TRAIN000000.JPG)
            * Final Label → 标签：
                "NRG" 表示 normal（无可转诊青光眼）
                "RG" 表示 Referable Glaucoma（可转诊青光眼）

    处理流程:
        1. 构造 Eye ID → 图像路径的映射。
        2. 解析 CSV，读取对应标签。
        3. 执行 crop + resize (512x512)。
        4. 保存到目标目录下的 images 文件夹。
        5. 输出 split.json 与 annotations.json。
        6. 输出统计信息到 ./experiments/stat/justraigs.json。

使用示例:
    python justraigs_preprocess.py \
        --data_path /home/zhangpinglu/data0/gy/Dataset/public_dataset/JustRAIGS \
        --tar_path /home/zhangpinglu/data0/gy/Dataset/public_processed/JustRAIGS \
        --prefix justraigs
"""

import os
import csv
import json
import argparse
from tqdm import tqdm
from data_preprocess.utils.crop import crop_resize_save
from data_preprocess.utils.stat import save_classification_stats


def parse_csv(csv_path):
    """
    读取 JustRAIGS_Train_labels.csv
    返回字典: {Eye ID: Final Label}
    """
    result = {}
    with open(csv_path, "r", encoding="utf-8-sig") as f:
        reader = csv.DictReader(f, delimiter=';')
        for row in reader:
            eye_id = row["Eye ID"].strip()
            label = row["Final Label"].strip()
            result[eye_id] = label
    return result


def build_id2path(data_path):
    """
    遍历 0~5 文件夹，构建 Eye ID -> 图片路径 映射。
    例如:
        TRAIN000000 -> /path/to/JustRAIGS/3/TRAIN000000.JPG
    """
    id2path = {}
    for subdir in map(str, range(6)):  # '0' ~ '5'
        folder = os.path.join(data_path, subdir)
        if not os.path.exists(folder):
            raise FileNotFoundError(f"子目录不存在: {folder}")
        for fname in os.listdir(folder):
            eye_id = os.path.splitext(fname)[0]
            abs_path = os.path.join(folder, fname)
            id2path[eye_id] = abs_path
    return id2path


def gather_data(data_path, tar_path, prefix="justraigs"):
    """
    主处理函数:
        1. 构建 id2path；
        2. 解析 CSV；
        3. 执行裁剪、resize、保存；
        4. 输出 annotations.json 和 split.json；
        5. 输出统计信息。
    """
    os.makedirs(tar_path, exist_ok=True)
    images_dir = os.path.join(tar_path, "images")
    os.makedirs(images_dir, exist_ok=True)

    csv_path = os.path.join(data_path, "JustRAIGS_Train_labels.csv")
    if not os.path.exists(csv_path):
        raise FileNotFoundError(f"CSV 文件不存在: {csv_path}")

    id2path = build_id2path(data_path)
    label_dict = parse_csv(csv_path)

    # 标签映射: Final Label -> 文本
    label_map = {
        "NRG": "normal",
        "RG": "Referable Glaucoma"
    }

    data_dict = {}
    split_dict = {"train": []}  # 当前数据集仅有 train 集

    print("\n🔹 开始处理 JustRAIGS 图像数据")
    for idx, (eye_id, label_code) in enumerate(tqdm(label_dict.items(), desc="processing", unit="images")):
        if eye_id not in id2path:
            raise FileNotFoundError(f"未找到对应图像: {eye_id} (CSV 中存在但文件缺失)")

        src_path = id2path[eye_id]
        if label_code not in label_map:
            raise ValueError(f"未知标签 '{label_code}' 于样本 {eye_id}")

        diagnosis_text = label_map[label_code]
        new_image_name = f"{prefix}_{eye_id}.png"
        dest_rel_path = os.path.join("images", new_image_name)
        dest_abs_path = os.path.join(tar_path, dest_rel_path)

        try:
            crop_info = crop_resize_save(
                image_path=src_path,
                save_path=dest_abs_path,
                resize=(512, 512),
                crop_threshold=25
            )
        except Exception as e:
            raise RuntimeError(f"处理图像失败: {src_path}, 错误信息: {str(e)}")

        data_dict[new_image_name] = {
            "image_name": new_image_name,
            "image_path": dest_abs_path,
            "original_path": src_path,
            "crop_info": crop_info,
            "diagnosis": {
                "classification": {"text": diagnosis_text}
            }
        }
        split_dict["train"].append(new_image_name)

    # 保存 split.json
    split_json_path = os.path.join(tar_path, "split.json")
    with open(split_json_path, "w", encoding="utf-8") as f:
        json.dump(split_dict, f, indent=4, ensure_ascii=False)

    # 保存 annotations.json
    annotations_json_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_json_path, "w", encoding="utf-8") as f:
        json.dump(data_dict, f, indent=4, ensure_ascii=False)

    # 保存分类统计结果
    os.makedirs("./experiments/stat", exist_ok=True)
    save_classification_stats(data_dict, os.path.join("./experiments/stat", "justraigs.json"))

    print("\n✅ 处理完成！共计:")
    print(f"  样本数: {len(split_dict['train'])} 张")
    print(f"  正常样本: {sum(1 for d in data_dict.values() if d['diagnosis']['classification']['text']=='normal')}")
    print(f"  可转诊样本: {sum(1 for d in data_dict.values() if d['diagnosis']['classification']['text']=='Referable Glaucoma')}")

    return data_dict


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="JustRAIGS 数据集预处理")
    parser.add_argument("--data_path", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_dataset/JustRAIGS",
                        help="原始数据集根目录，包含 0~5 文件夹和 JustRAIGS_Train_labels.csv")
    parser.add_argument("--tar_path", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_processed/JustRAIGS",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="justraigs",
                        help="输出图片文件名前缀")
    args = parser.parse_args()

    gather_data(args.data_path, args.tar_path, prefix=args.prefix)
